from paddle.fluid.dygraph import Layer, Sequential, Dropout

from helm.static.models.layers import Conv2d, GlobalAvgPool, Linear, Identity, LayerList

__all__ = [
    'MobileNetV2'
]



class InvertedResidual(Layer):
    def __init__(self, in_channels, out_channels, channels, t=6, strides=1):
        super(InvertedResidual, self).__init__()
        self.strides = strides
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.t = t

        self.bottleneck = Sequential(
            Conv2d(in_channels, channels * t, kernel_size=1, bn=True, act='relu6'),
            Conv2d(channels * t, channels * t, kernel_size=3, bn=True, act='relu6', stride=strides, groups=channels * t),
            Conv2d(channels * t, out_channels, kernel_size=1, bn=True, act=None),
        )

        # if in_channels != out_channels and strides != 2:
        #     self.shortcut = Conv2d(in_channels, out_channels, kernel_size=1, bn=True, act=None)
        # else:
        #     self.shortcut = Identity()

    def forward(self, x):
        identity = x
        x = self.bottleneck(x)
        if self.in_channels == self.out_channels and self.strides == 1:
            x = x + identity
        return x


# class Bottleneck(Layer):
#     def __init__(self, num_filters_in, num_filters_out, t=6, n=1, strides=1):
#         super(Bottleneck, self).__init__()
#         self.strides = strides
#         self.num_filters_in = num_filters_in
#         self.num_filters_out = num_filters_out
#         self.t = t
#         self.n = n
#
#         self.bottleneck = LayerList()
#
#         self.bottleneck.append(
#             InvertedResidual(num_filters_in, num_filters_out, num_filters_in, t, strides)
#         )
#         for i in range(1, n):
#             self.bottleneck.append(
#                 InvertedResidual(num_filters_out, num_filters_out, num_filters_in, t, 1)
#             )
#
#     def forward(self, x):
#         for i in range(self.n):
#             x = self.bottleneck[i](x)
#         return x


class MobileNetV2(Layer):
    def __init__(self, width_multiplier=1.0, num_classes=10):
        super(MobileNetV2, self).__init__()
        channel_list = [int(round(x * width_multiplier, 2)) for x in [32, 16, 24, 32, 64, 96, 160, 320, 1280]]

        self.conv1 = Conv2d(3, channel_list[0], kernel_size=3, bn=True, act='relu6')

        self.bottleneck_layer = LayerList()
        param = [[channel_list[0], 6, channel_list[1], 1, 1],
                 [channel_list[1], 6, channel_list[2], 2, 1],
                 [channel_list[2], 6, channel_list[3], 3, 2],
                 [channel_list[3], 6, channel_list[4], 4, 2],
                 [channel_list[4], 6, channel_list[5], 3, 1],
                 [channel_list[5], 6, channel_list[6], 3, 2],
                 [channel_list[6], 6, channel_list[7], 1, 1],
                ]
        for c_in, t, c_out, n, s in param:
            self.bottleneck_layer.append(self._make_layer(c_in, c_out, t, n, s))

        self.bottlenecks = Sequential(*self.bottleneck_layer)
        self.conv2 = Conv2d(channel_list[7], channel_list[8], kernel_size=1, bn=True, act='relu6')
        self.pool = GlobalAvgPool()
        self.dense = Linear(channel_list[8], num_classes)

    def _make_layer(self, in_channels, out_channels, t=6, n=1, strides=1):
        layers = LayerList()
        layers.append(
            InvertedResidual(in_channels, out_channels, in_channels, t, strides)
        )
        for i in range(1, n):
            layers.append(
                InvertedResidual(out_channels, out_channels, in_channels, t, 1)
            )
        inverted_residual = Sequential(*layers)
        return inverted_residual


    def forward(self, x):
        x = self.conv1(x)
        x = self.bottlenecks(x)
        x = self.conv2(x)
        x = self.pool(x)
        x = self.dense(x)
        return x
